import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms.functional
from torchvision import transforms
import numpy as np


class torch_nn_functional_usage_class(nn.Module):
    """
    定义几个模型训练中常用的函数，并演示使用效果
    """

    def __init__(self):
        super(torch_nn_functional_usage_class, self).__init__()
        print("* 注意:需要引入的是 \033[1;31m`import torch.nn.functional as F` \033[0m")


class torch_nn_forward_usage_class(nn.Module):
    """
    通过学习了nn.layer和nn.functional部分，现在已经可以对一个较为简单的网络进行学习了
    下面将通过简单实例来展示一个完整的模型的代码。
    """

    def __init__(self):
        super(torch_nn_forward_usage_class, self).__init__()
        print("\033[1;31m——————————————————————简单完整模型使用实例—————————————————————————\033[0m")

        self.conv1 = nn.Conv2d(3, 6, 5, padding=2)
        self.pool1 = nn.MaxPool2d(2, 2, padding=0)
        self.conv2 = nn.Conv2d(6, 16, 5, padding=2)
        self.pool2 = nn.MaxPool2d(2, 2, padding=0)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)
        print("该模型定义了多个层：\n")
        print("conv1: 输入通道为{0}, 输出通道为{1}, 核大小为 {2} "
              .format(self.conv1.in_channels, self.conv1.out_channels, self.conv1.kernel_size, ))
        print("pool1: 核大小为{0}, 步长为{1} "
              .format(self.pool1.kernel_size, self.pool1.stride))
        print("conv2: 输入通道为{0}, 输出通道为{1}, 核大小为 {2} "
              .format(self.conv2.in_channels, self.conv2.out_channels, self.conv2.kernel_size, ))
        print("pool2: 核大小为{0}, 步长为{1} "
              .format(self.pool2.kernel_size, self.pool2.stride))
        print("fc1: 输入通道为{0}, 输出通道为{1}"
              .format(self.fc1.in_features, self.fc1.out_features))
        print("fc2: 输入通道为{0}, 输出通道为{1}"
              .format(self.fc2.in_features, self.fc2.out_features))
        print("fc3: 输入通道为{0}, 输出通道为{1}"
              .format(self.fc3.in_features, self.fc3.out_features))
        print("\033[1;33m当这些组件定义好之后，就可以定义 forward()函数，用来搭建网络结构\033[0m")

    def forward(self, x):
        """
        :param x: 模型的输入
        :return:
        """
        size0 = x.size
        print("图像原尺寸为：{0}".format(size0))

        # 尺寸变换
        print("首先将图像缩放到指定尺寸：1920×1080")
        x = torchvision.transforms.functional.resize(x, [1920, 1080])

        # 类型转换
        print("通过使用实例 fun = transforms.ToTensor()作为函数, 将图片转为张量")
        toTensor = transforms.ToTensor()
        x = toTensor(x)
        print("张量X的输入尺寸为：{0}".format(x.size()))
        image_size = x.size()
        h = image_size[1]
        w = image_size[2]
        print("-----------------------------------------------------")
        print("\033[1;33m转为张量后，其尺寸为{0}，需要转变为4维张量，将张量作为一个整体，嵌到一个维度里即可\033[0m".format(image_size))
        print("通过 X = X.reshape((1, 3, h, w))来实现")
        x = x.reshape((1, 3, h, w))
        size1 = x.size()
        print("现在 X 的尺寸为{0}".format(x.size()))

        # 卷积和池化
        x = self.conv1(x)
        size2 = x.size()
        x = self.pool1(F.relu(x))
        size3 = x.size()
        x = self.conv2(x)
        size4 = x.size()
        x = self.pool2(F.relu(x))
        size5 = x.size()

        # 这里进行重采样
        x = x.view(-1, 16 * 5 * 5)

        # 进行线性变换
        x = F.relu(self.fc1(x))
        size6 = x.size()
        x = F.relu(self.fc2(x))
        size7 = x.size()
        x = self.fc3(x)
        size8 = x.size()
        print("1. x 经过 conv1，然后经过激活函数 relu，再经过 pool1 操作")
        print("2. x 经过 conv2，然后经过激活函数 relu，再经过 pool2 操作")
        print("3. 将 x 进行 reshape，为了后面做为全连接层的输入")
        print("4. 经过全连接层 fc1，然后经过 relu")
        print("5. 经过全连接层 fc2，然后经过 relu")
        print("6. 经过全连接层 fc3")
        print("-----------------------------------------------------")
        print("\033[1;34m尺寸变化示意\033[0m：")
        print("\t原始图像大小：\t{0}".format(size0))
        print("\t图像转为张量：\t{0}".format(size1))
        print("\t                      \t↓")
        print("\t\033[1;31mconv1卷积\033[0m：\t{0}".format(size2))
        print("\t\033[1;31mpool1池化\033[0m：\t{0}".format(size3))
        print("\t\033[1;31mconv2卷积\033[0m：\t{0}".format(size4))
        print("\t\033[1;31mpool2池化\033[0m：\t{0}".format(size5))
        print("\t\033[1;32m这里进行了reshape\033[0m      \t↓")
        print("\t\033[1;31mfc1线性变换\033[0m：\t{0}".format(size6))
        print("\t\033[1;31mfc2线性变换\033[0m：\t{0}".format(size7))
        print("\t\033[1;31mfc3线性变换\033[0m：\t{0}".format(size8))


class torch_nn_backward_usage_class(nn.Module):
    """
    这是一个有反向传播的模型的使用示例
    通过ResNe34作为例子使用
    """


class ResNet34(nn.Module):
    """
    这是ResNet32模型的完整复现，用来学习深度学习的基本框架
    * ResNet34 包含多个 layer，每个 layer 又包含多个 Residual block
    模型的网络架构可以在论文中找到。

    #该模型建立的时候用到了几个技巧，
    ① torch.nn.Sequential 其实就是 Sequential 容器，该容器将一系列操作按先后顺序给包起 来方便重复使用，
    ② 创建子Module，用作批量创建堆叠层，

    小结：通过本次模型的创建，有以下需要注意的：
    ——模型的定义就是先继承，再构建组件，最后组装。




    # 初始化权值的方法  在函数 initialize_weights 中实现
    通过调用 self.modules()可以获取该模型中的所有层级，
    通过调用 isinstance(m, Instance) 可以判断 m 是否为 Instance 类型
    模型的初始化方法一般可以调用 torch.nn.init 下的方法来进行，具体可以在 initialize_weights 中查看


    """

    def __init__(self):
        super(ResNet34, self).__init__()
        self.model_name = 'ResNet34'
        """
        输入image 224x224x3
                7x7 conv, 64, /2
        3个【小块】 3x3 conv, 64 
                3x3 conv, 128, /2
        4个【小块】 3x3 conv, 128 
                3x3 conv, 256, /2
        6个【小块】 3x3 conv, 256
                3x3 conv, 512, /2
        3个【小块】 3x3 conv, 512
        一个fc1000的avg pool
        
        模型中不仅有直连层，而且还分大小块，每两层作为一个小块，
        """
        self.origin_W = 224
        self.origin_H = 224
        self.C0 = 3  # 输入是RGB色彩图片，也就是三通道
        self.k0 = 7

        self.C1 = 64
        self.k1 = 3

        self.C2 = 128
        self.k2 = 3

        self.C3 = 256
        self.k3 = 3

        self.C4 = 512
        self.k4 = 3

        self.Ce = 1000  # 分类的数量

        # 定义预处理网络 7x7 conv, 64, /2
        self.pre = nn.Sequential(
            nn.Conv2d(self.C0, self.C1, self.k0, stride=2, padding=int((self.k0 - 1) / 2), bias=False),
            nn.BatchNorm2d(self.C1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1))
        # 接下来是重复构造出来的几个层级，写一个函数批量定义
        self.layer1 = self._make_layer(self.C1, self.C2, 3, do_half='no')
        self.layer2 = self._make_layer(self.C2, self.C3, 4, do_half='yes')
        self.layer3 = self._make_layer(self.C3, self.C4, 6, do_half='yes')
        self.layer4 = self._make_layer(self.C4, self.C4, 3, do_half='yes')
        self.fc = nn.Linear(self.C4, self.Ce)

    def _make_layer(self, Ci, Co, block_num, do_half='no'):
        """
        批量构建残差堆叠层residual block

        """
        # 这个是shortcut层，由于每一个小的堆叠层有两个卷积层，和一个直连层，
        # 为了方便，需要用到子模型来批量构建，因此创建一个新的模型类class ResidualBlock(nn.Module)

        stride = 1 if do_half == 'no' else 2  # 判断该部分的开头是否要减半

        shortcut = nn.Sequential(nn.Conv2d(Ci, Co, 1, stride, bias=False),
                                 nn.BatchNorm2d(Co))
        # 每一部分开头传入的是减半后的特征图，而之后的不是，为了传入残差，这里需要单独传入一个减半后的特征图
        layers = [ResidualBlock(Ci, Co, stride, shortcut)]
        # 构造之后的 block_num - 1 个残差块
        for i in range(1, block_num):
            layers.append(ResidualBlock(Co, Co))  # 后面是同输入同输出的

        return nn.Sequential(*layers)

    def forward(self, x):

        x = torchvision.transforms.functional.resize(x, [224, 224])
        x = transforms.ToTensor()(x)
        x = x.reshape((1, 3, self.origin_H, self.origin_W))  # 这一步总是非常重要的，不然不能运行
        x = self.pre(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = F.avg_pool2d(x, 7)
        x = x.view(x.size(0), -1)  # 这里进行重新组装 改造成线性的，然后给出1000中样本分类的结果
        x = self.fc(x)
        return x

    # 接下来要重复组装子块，

    def get_FLOPs(self):
        """
        TODO: 该部分还没有完成
        """
        print("作者给出的ResNet的FLOPs为 (3.6 billion FLOPs)")
        print("现给出公式运算并检验计算结果")
        """

        """

        FLOPs = 0

        W = self.origin_W
        H = self.origin_H
        # 预处理部分
        FLOPs = FLOPs + 2 * W * H * self.C0 * self.C1 * self.k0 ** 2
        # 第一部分的 6 次卷积
        W, H = W / 2, H / 2
        FLOPs = FLOPs + 2 * W * H * self.C1 * self.C2 * self.k1 ** 2 * 6
        # 第二部分的 7 次卷积
        W, H = W / 2, H / 2
        FLOPs = FLOPs + 2 * W * H * self.C2 * self.C3 * self.k2 ** 2 * 7
        # 第三部分的 11 次卷积
        W, H = W / 2, H / 2
        FLOPs = FLOPs + 2 * W * H * self.C3 * self.C4 * self.k3 ** 2
        # 第四部分的 5 次卷积
        W, H = W / 2, H / 2
        FLOPs = FLOPs + 2 * W * H * self.C4 * self.C4 * self.k4 ** 2
        # 最后一部分的输出 不明确到底多少
        FLOPs = FLOPs + W * H * self.C4 * self.Ce

        print("模型的FLOPs为{0} G".format(FLOPs / 1000 ** 3))

    def initialize_weights(self):
        """
        这段代码基本流程是这样，先从 self.modules()中遍历每一层，然后判断各层属于什么
        类型，例如，是否是 nn.Conv2d、nn.BatchNorm2d、nn.Linear 等，然后根据不同类型的
        层，设定不同的权值初始化方法，例如，Xavier，kaiming，normal_，uniform_等。
        Ps: kaiming 也称之为 MSRA 初始化
        """

        # self.modules() Returns an iterator over all modules in the network. 能依次返回模型中的各层

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                torch.nn.init.xavier_normal(m.weight.data)
                if m.bias is not None:
                    m.bias.data.zero_()  # 偏置项置为零
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)  # 置为1
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                torch.nn.init.normal(m.weight.data, 0, 0.01)
                m.bias.data.zero_()


class ResidualBlock(nn.Module):
    """
    残差块由两个层组成，基本构成是
    conv1 - batch1 - relu1
    conv2 - batch2 - relu2

    """

    def __init__(self, Ci, Co, stride=1, shortcut=None):
        super(ResidualBlock, self).__init__()
        self.left = nn.Sequential(nn.Conv2d(Ci, Co, 3, stride, 1, bias=False),  # 3 × 3 卷积部分，由于可能是第一个开头的增倍层，所以设为参数stride
                                  nn.BatchNorm2d(Co),
                                  nn.ReLU(inplace=True),  # 激活函数
                                  nn.Conv2d(Co, Co, 3, 1, 1, bias=False),  # 3 × 3 卷积部分
                                  nn.BatchNorm2d(Co))
        # 这个是连接层，可能有可能没有
        self.right = shortcut

    def forward(self, x):
        out = self.left(x)
        residual = x if self.right is None else self.right(x)  # 有些块的残差经过了线性变换（每部分的第一个块）
        out += residual
        return F.relu(out)
